#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""分拣模拟系统API接口模块"""
from fastapi import APIRouter, Query, HTTPException, WebSocket, WebSocketDisconnect
from typing import Dict, Any
from fastapiApp.utils.luosi_run import get_simulation
from fastapiApp.utils.sorting import get_sorting_simulator
import json
import asyncio

router = APIRouter()

# WebSocket连接管理器
class ConnectionManager:
    """WebSocket连接管理器，用于管理所有客户端连接"""

    def __init__(self):
        # 存储所有活跃连接
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        """接受WebSocket连接并添加到活跃连接列表"""
        await websocket.accept()
        self.active_connections.append(websocket)
        print(f"新的WebSocket连接已建立，当前连接数: {len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket):
        """从活跃连接列表中移除断开的连接"""
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
            print(f"WebSocket连接已断开，当前连接数: {len(self.active_connections)}")

    async def send_personal_message(self, message: str, websocket: WebSocket):
        """向特定客户端发送消息"""
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        """向所有连接的客户端广播消息"""
        for connection in self.active_connections:
            try:
                await connection.send_text(message)
            except Exception as e:
                print(f"发送消息失败: {e}")
                # 如果发送失败，从连接列表中移除该连接
                self.active_connections.remove(connection)

# 创建全局连接管理器实例
manager = ConnectionManager()

def get_simulation_instance():
    """获取模拟实例"""
    return get_simulation()

def get_sorting_simulator_instance():
    """获取分拣模拟系统实例"""
    return get_sorting_simulator()

@router.get("/start", summary="启动模拟")
async def start_simulation():
    """启动分拣模拟系统"""
    simulator = get_sorting_simulator_instance()
    await simulator.start_simulation()
    return "start success"

@router.get("/stop", summary="停止模拟")
async def stop_simulation():
    """停止分拣模拟系统"""
    simulator = get_sorting_simulator_instance()
    await simulator.stop_simulation()
    return "stop success"

@router.get("/init", summary="初始化模拟")
async def init_simulation():
    """初始化分拣模拟系统"""
    simulator = get_sorting_simulator_instance()
    await simulator.init_simulation()
    return "init success"

@router.get("/getstatus", summary="获取模拟状态")
async def get_simulation_status():
    """获取分拣模拟系统状态"""
    simulator = get_sorting_simulator_instance()
    return simulator.get_state()

@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket端点，用于实时推送模拟状态"""
    await manager.connect(websocket)
    try:
        # 首次连接时发送当前状态
        simulator = get_sorting_simulator_instance()
        current_state = simulator.get_state()
        await manager.send_personal_message(json.dumps(current_state), websocket)

        # 持续监听客户端消息
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            message_data = json.loads(data)

            # 处理不同类型的消息
            if message_data.get("type") == "get_status":
                # 客户端请求获取最新状态
                simulator = get_sorting_simulator_instance()
                state = simulator.get_state()
                await manager.send_personal_message(json.dumps(state), websocket)

            elif message_data.get("type") == "start":
                # 客户端请求启动模拟
                await start_simulation()
                # 发送更新后的状态
                simulator = get_sorting_simulator_instance()
                state = simulator.get_state()
                await manager.send_personal_message(json.dumps(state), websocket)

            elif message_data.get("type") == "stop":
                # 客户端请求停止模拟
                await stop_simulation()
                # 发送更新后的状态
                simulator = get_sorting_simulator_instance()
                state = simulator.get_state()
                await manager.send_personal_message(json.dumps(state), websocket)

            elif message_data.get("type") == "init":
                # 客户端请求初始化模拟
                await init_simulation()
                print("初始化模拟系统")
                # 发送更新后的状态
                simulator = get_sorting_simulator_instance()
                state = simulator.get_state()
                await manager.send_personal_message(json.dumps(state), websocket)

    except WebSocketDisconnect:
        manager.disconnect(websocket)
        print("WebSocket连接已断开")
    except Exception as e:
        print(f"WebSocket错误: {e}")
        manager.disconnect(websocket)

# 后台任务，用于定期广播状态更新
async def broadcast_state_updates():
    """定期广播状态更新给所有连接的客户端"""
    while True:
        if manager.active_connections:
            simulator = get_sorting_simulator_instance()
            state = simulator.get_state()
            await manager.broadcast(json.dumps(state))
        # 每秒更新一次
        await asyncio.sleep(1)



@router.get("/init_simulation_one", response_model=Dict[str, Any], summary="初始化模拟系统")
async def init_simulation_one():
    """初始化分拣模拟系统"""
    simulation = get_simulation_instance()
    simulation.init_simulation()
    return simulation.get_state()

@router.get("/state", response_model=Dict[str, Any], summary="获取当前状态")
async def get_simulation_state():
    """获取分拣模拟系统的当前状态"""
    simulation = get_simulation_instance()
    return simulation.get_state()

@router.post("/reset", response_model=Dict[str, Any], summary="重置模拟")
async def reset_simulation():
    """重置分拣模拟系统"""
    simulation = get_simulation_instance()
    return simulation.reset_simulation()

@router.post("/toggle-auto-run", response_model=Dict[str, Any], summary="切换自动运行")
async def toggle_auto_run():
    """切换自动运行状态"""
    simulation = get_simulation_instance()
    return simulation.toggle_auto_run()

@router.post("/step-run", response_model=Dict[str, Any], summary="单步运行")
async def step_run():
    """单步运行模拟"""
    simulation = get_simulation_instance()
    return simulation.step_run()

@router.post("/change-station-size", response_model=Dict[str, Any], summary="改变中转站大小")
async def change_station_size(size: int = Query(..., ge=6, le=30, description="中转站库位数量（6-30之间）")):
    """改变中转站库位数量"""
    simulation = get_simulation_instance()
    return simulation.change_station_size(size)

@router.post("/change-product-types-count", response_model=Dict[str, Any], summary="改变产品种类数量")
async def change_product_types_count(count: int = Query(..., ge=2, le=20, description="产品种类数量（2-20之间）")):
    """改变产品种类数量"""
    simulation = get_simulation_instance()
    return simulation.change_product_types_count(count)

@router.get("/product-stats", response_model=Dict[str, Any], summary="获取产品统计")
async def get_product_stats():
    """获取产品统计信息"""
    simulation = get_simulation_instance()
    state = simulation.get_state()
    return state["product_stats"]

@router.get("/logs", response_model=Dict[str, Any], summary="获取日志")
async def get_logs():
    """获取操作日志"""
    simulation = get_simulation_instance()
    state = simulation.get_state()
    return {"logs": state["logs"]}

@router.post("/toggle-product-generation-mode", response_model=Dict[str, Any], summary="切换产品生成模式")
async def toggle_product_generation_mode():
    """切换产品生成模式（随机生成 <-> 顺序生成）"""
    simulation = get_simulation_instance()
    return simulation.toggle_product_generation_mode()
